// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//

#include <kllm/openai/tokenize.h>

namespace kllm {

    turbo::Status format_aoi_tokenize_response(const KaiResponse &response,  KMContext *ctx, nlohmann::ordered_json &json_res) {
        auto with_pieces = response.with_pieces();
        auto &tokens = response.tokens();
        auto obj = nlohmann::ordered_json::array();
        if(with_pieces) {
            for (const auto& token : tokens) {
                std::string piece = ctx->token_to_piece(token);
                nlohmann::ordered_json piece_json;

                // Check if the piece is valid UTF-8
                if (is_valid_utf8(piece)) {
                    piece_json = piece;
                } else {
                    // If not valid UTF-8, store as array of byte values
                    piece_json = nlohmann::ordered_json::array();
                    for (unsigned char c : piece) {
                        piece_json.push_back(static_cast<int>(c));
                    }
                }

                obj.push_back({
                                                  {"id", token},
                                                  {"piece", piece_json}
                                          });
            }
        } else {
            obj = tokens;
        }
        json_res["tokens"] = obj;
        return turbo::OkStatus();
    }

    turbo::Status format_aoi_detokenize_response(const KaiResponse &response, KMContext *ctx, nlohmann::ordered_json &obj) {
        obj =  nlohmann::ordered_json{
                {"content", response.content()}
        };
        return turbo::OkStatus();
    }
}  // namespace kllm